#!/usr/bin/env python3
#coding=utf-8

import string
import random
import base64
import _thread # for compatability with python2
import time
import datetime
import sys

import logging
import coloredlogs

import pymongo
from bson import binary
from bson.son import SON 
from bson.code import Code
from pymongo import ReadPreference
from pymongo.mongo_client import MongoClient

# https://juejin.im/post/6844903597465927694

host=''

uri_localhost='mongodb://localhost:27017'
uri_local_rs='mongodb://admin:123456@localhost:30001'

logger = logging.getLogger('main')

coloredlogs.install(level='DEBUG')

# https://stackoverflow.com/questions/2257441/random-string-generation-with-upper-case-letters-and-digits
def random_string(n, chars=string.ascii_lowercase + string.ascii_uppercase + string.digits):
    return ''.join(random.choice(chars) for _ in range(n))

def insert_data(col, n):
    for i in range(n):
        b64 = base64.b64encode(random_string(16))
        item = {'k': i, 'v1': binary.Binary(b64)}
        col.insert(item)

def find_data(col):
    pass

# 如何并发校验 ？？每个连接线程约定好各自取不同段的数据 ？？
def check_data():
    c = pymongo.MongoClient(rocks_uri, readPreference='primary')
    c2 = pymongo.MongoClient(rocks_secondary_uri, readPreference='secondary')
    t = c[db][col]
    res = t.find()
    for doc in res:
        id = doc['_id']
        #print 'id: ', id
        slave_doc = c2[db][col].find_one({'_id': id})
        if doc != slave_doc:
            print('doc=', doc)
            print('slave_doc', slave_doc)
            print('ERROR !')
            sys.exit(-1)
        time.sleep(2)

def get_col():
    pass

def thread_show_col_names(i, uri):
    client = pymongo.MongoClient(uri)
    db = client.db1
    print('thread: ', i)
    while True:
        col = db.collection_names()
        for c in col:
            stats = db.command('collstats', c)
            line = '{},{},{},{}'.format(stats['ns'], stats['size'],stats['count'], stats['avgObjSize'],
                stats['storageSize'], stats['nindexes'], stats['totalIndexSize'])
            print(line)
        
        time.sleep(3)
    print('thread: ', i, ' exited')

def t_launch_threads():
    for i in range(1,2):
        _thread.start_new_thread(thread_show_col_names, (i, uri_local_rs))

    time.sleep(3600)

def t_get_all_col_stats(c):
    f = open('all-col.csv', 'w')
    f.write('ns,size,count,avgObjSize,storageSize,nindexes,totalIndexSize\n')
    dblist = c.database_names()
    for dbname in dblist:
        if dbname == 'local' or dbname == 'admin' or dbname == 'config':
            continue

        db = c[dbname]
        col = db.collection_names()
        for c in col:
            stats = db.command('collstats', c)
            line = '{},{},{},{},{},{},{}'.format(stats['ns'], stats['size'],stats['count'], stats['avgObjSize'],
                stats['storageSize'], stats['nindexes'], stats['totalIndexSize'])
            f.write(line + '\n')

    f.close()

def t_create_many_collections(c):
    start = 10000
    len = 30000
    db = c['db1']
    for i in range(start, len):
        begin = datetime.datetime.now()
        a = time.time()
        db['col_' + str(i)].create_index([('v1', pymongo.DESCENDING)], background=True)

        for j in range(100):
            db['col_' + str(i)].insert({'k': i+j, 'v1': i+j+1, 'v2': i+j+2, 'v3': i+j+3})

        db['col_' + str(i)].create_index('v2', background=True)
        end = datetime.datetime.now()
        b = time.time()

        print("insert {} took {}, {}".format(i, (end-begin).microseconds / 1000, (b-a)))

def t_use_secondaryPreferred(mongos_list, colname):
    for i in range(len(mongos_list)):
        uri = 'mongodb://admin:123456@' + mongos_list[i] + '/admin?readPreference=secondaryPreferred'
        #print('uri:', uri)
        c = MongoClient(uri)
        col = c['db2'][colname]
        for item in col.find({'k': 1}):
            print('mongos:', mongos_list[i],  item)

def tt(colname):
    uri = "mongodb://admin:123456@9.134.5.74:30000/admin?readPreference=secondaryPreferred"                                                     
    print("read from ", uri)                                                                                                                                                                                                                                                                                                                             
                                                                                                                                                                                    
    x = MongoClient(uri)                                                                                                                                                             
    db = x.db2                                                                                                                                                         
    co = db[colname]                                                                                                                                                                   
    for item in co.find({"k":1}):                                                                                                                                                
        print(item, "\n")                                                                                                                                                               
                                                                                                                                                                                    
    uri = "mongodb://admin:123456@9.134.5.74:40000/admin?readPreference=secondaryPreferred"                                                     
    print("read from ", uri)                                                                                                                                                         
    x = MongoClient(uri)                                                                                                                                                             
    db = x.db2                                                                                                                                                         
    co = db[colname]                                                                                                                                                                   
    for item in co.find({"k":1}):                                                                                                                                                
        print(item, "\n")    

if __name__ == "__main__":
    #t_use_secondaryPreferred(sys.argv[1].split(','), sys.argv[2])
    tt(sys.argv[1])

    #uri = sys.argv[1]
    #c = pymongo.MongoClient(host=uri)

    # 貌似下面这个uri就不能指定主？否则连的还是主？？
    # c2 = pymongo.MongoClient(rocks_secondary_uri, readPreference='secondary')
    # how to print client ip ????????
    #print c.address
    #print random_string(16)
    #insert_data(c.db1.t, 10)

    #t_create_many_collections(c)

    #t_launch_threads()

    #t_get_all_col_stats(c)
    print('done')
